package beast.phylodynamics.marsden;

import beast.evolution.tree.Tree;
import beast.evolution.tree.coalescent.TreeIntervals;
import beast.phylodynamics.epidemiology.DeterministicCoalescentSIR;
import beast.util.TreeParser;
import beast.core.parameter.RealParameter;
import beast.phylodynamics.epidemiology.Volz2009TreeDistribution;

// Restructured since Denise wrote this -Alex

//import beast.phylodynamics.BDSIR;

import java.io.IOException;
import java.io.PrintStream;

/**
 * Created by IntelliJ IDEA. User: Denise Kuehnert Date: Jun 4, 2013 Time:
 * 4:36:41 PM
 */
public class LikelihoodCrossSections {

    public static void main(String [] args) throws IOException, Exception {

        PrintStream writer;
        
        String [] param_name = {"beta", "gamma", "S0"};
        double[] paramTruth = {0.0075, 0.30, 999.0};
        double[] paramLowerBound = {0.0001, 0.10, 500.0};
        double[] paramUpperBound = {0.01, 0.50, 1500.0};
        int nsteps = 100;
        
        double [] paramCurrent = new double[3];

        String[] trees = new String[1];
        int[] tips = new int[trees.length];
        
        //25 tips
//        trees[0] = "((20:1.7327013309920987,(2:1.5699517642369987,9:1.5699517642369987):0.1627495667551):0.3906420437073396,(((((12:1.4606660366895723,((22:0.7364772016601255,21:0.7364772016601255):0.6701138844558105,(15:1.3665082235172927,((3:1.020945860930602,(8:0.8653300121493379,10:0.8653300121493379):0.155615848781264):0.22753196406186138,(7:1.1290942502849537,19:1.1290942502849537):0.1193835747075096):0.11803039852482944):0.040082862598643354):0.054074950573636205):0.0074844678674776954,((17:0.012266430667138728,23:0.012266430667138728):1.3007893648454143,(24:0.9390090648133209,5:0.9390090648133209):0.3740467306992321):0.15509470904449696):0.038425441676050665,((18:0.5287382111016794,16:0.5287382111016794):0.6219873839148018,14:1.1507255950164812):0.35585035121661945):0.17026942978592519,11:1.6768453760190258):0.3913503628622723,((13:0.7724098713155667,(25:0.07463738743032255,4:0.07463738743032255):0.6977724838852442):0.9249671690996732,(6:1.2765864559055629,1:1.2765864559055629):0.4207905845096771):0.3708186984660582):0.05514763581814014);";
//        tips[0] = 25;

        //100 tips
        trees[0] = "((((30:1.8551845598088046,(84:0.1426246794287902,(((43:0.9287981039828717,13:0.0758872076210908):0.04946986792500052,(((((68:0.05607068314827579,(90:0.24258318089350617,72:0.8560880915569644):0.5372180006395109):1.0858610536359183,17:1.6263791239132104):1.1111352073370764,33:0.38964229106509407):0.04547746758987614,(((7:0.5867835761176927,58:1.546103110785534):0.8335743279618999,61:2.1355856044173507):0.17691191534672512,77:1.062855596386063):1.214190850577559):3.7957901293609275,((95:1.6383475772471252,54:0.28737859379128405):0.013409641527529104,(69:0.42796923487643657,(((((((16:1.8094355177088843,99:1.4420385635861521):0.823908663553139,(2:0.8558884912972449,(97:0.6367547383210592,(23:0.16686753940627597,79:0.8324930978707314):0.6215625238860323):0.260292285075133):1.2909447556209344):1.9403525431395652,(45:0.18000771309106867,26:0.1646382885369757):2.2941500862976065):0.29246569617450824,(78:0.23334706328929222,((((91:0.20052573318247013,44:0.36372247188358386):0.13999837788867886,(((11:1.482973483181972,29:1.8302755671138158):0.9126110424822755,38:0.5127382350568954):0.005981948312886232,89:0.35779423124706256):0.2326595640973217):0.09576407740722459,(40:1.139261786541951,1:0.22018445937417752):0.17733923495773674):0.46602893061873196,18:1.6763252778644926):0.9562179850124952):0.9049894880743237):0.40618035645014317,(((75:2.7653028821638754,((76:1.6902625931990745,(10:0.2969523066575519,((62:1.5758226231093193,5:0.4969597314518417):0.13328791654158856,83:0.2177750674321306):0.2131929723353032):0.29031481448532404):1.0013382351231321,93:2.55308929905547):0.9324768285610068):0.07625892830646208,((25:0.03233815767178072,(57:0.5173815091349585,((((63:1.1432711046929338,22:1.2409817802001655):0.5206526876791493,81:0.7227138118634979):0.22946012771625313,51:0.9132059655628151):0.3325049397909652,92:2.2138384688941652):0.36100550228978534):1.6621960185959352):0.567318495911266,((74:1.6489350707337742,41:1.004742320939462):0.37940330214507156,(((86:0.7484146773743703,(24:1.7948836089067832,(28:0.10744680009377738,(15:0.5976788684992815,50:0.9573191410800241):0.7737872462363651):0.23125408934417813):0.39952099502231775):0.9589673782168369,((64:1.369119021143387,(94:0.0916778479809306,(71:0.08057497851065598,98:0.020960408408679854):0.23142977466553383):0.23474015908246848):0.33791173372171635,(100:0.1421067485789198,14:0.2574518586515211):1.838212867605863):1.860179332329336):0.01989568483846682,((21:0.0658070338372383,(31:1.227937810696993,(4:0.9936427013471469,(96:0.4227926209119044,48:1.0814462185076987):1.4658848661551396):0.2286731129736701):0.4661095972034097):0.3680656177197754,((39:0.2380909048460751,(60:0.84240292031091,9:0.7504474998134807):0.4597316544074861):1.1379507631011112,(53:0.8563414173405768,35:0.5222915571431486):1.8707557863546729):0.19298013054743102):0.6757651500916602):0.31067579455440075):0.20589360386820843):0.6688629062170426):0.05797729409275121,(80:2.0976108053303895,42:0.49006520631621964):0.5981075262451601):0.18058426167347186):0.16494610913142793,19:4.743101654252418):0.40972007758361784,(27:0.30093370164305533,37:0.18367898024736284):2.6372657248090556):0.07473203909910975):1.109866122795144):0.07889204895055268):0.08167068377693809):0.08141267443441702,55:0.3505922364617575):1.106083003110479):0.7270388997958848):1.0859220630399213,73:0.39867523449442954):0.8489706299731228,(((52:0.09398645231272695,(70:1.60685294521784,85:1.815107670823819):1.3853146629295425):3.2511246569221806,((((((56:0.5994269084458193,3:1.4738729532730037):0.803892885870571,34:1.238301698850675):2.2134490414381904,6:0.32473549565563964):1.0561539296475573,(88:0.28667280172115994,66:1.2081728239099032):0.4457411149358812):0.5370895580396597,((((46:0.07975636623612203,(49:1.428647986503238,32:0.9319426123447823):2.085187259028972):0.7245855147803395,47:0.0808854545852764):0.08708882381464633,82:1.3885174683227453):7.674560520420926E-6,((36:1.2915030129307024,67:0.6707440615262783):0.9490411510584487,(59:1.0337282065508973,(65:0.6709567708862618,8:0.8490105063522542):2.081022256435972):0.1346838163447881):1.6709324015080078):0.4665379007974879):0.9864768029684434,12:0.9538850850499196):1.2853050159777153):0.11673127884672763,20:0.6370276607175409):2.936788560310112):0.07428776159986139,87:1.0318509846459885):1.3501679528862376";
        tips[0] = 100;

        // all 245 tips
//        trees[2] = "(((((((((((233:0.6955967111603156,198:0.6955967111603156):0.14021726278152546,79:0.8358139739418411):0.2371408222226441,(172:0.35672813225298183,(30:0.11114029483538612,35:0.11114029483538612):0.2455878374175957):0.7162266639115034):0.096868837962653,(((38:0.3735191185528217,32:0.3735191185528217):0.39229959315571605,(224:0.3886510272331738,229:0.3886510272331738):0.377167684475364):0.08685215563597892,(101:0.24682061936594568,196:0.24682061936594568):0.605850247978571):0.3171527667826215):0.14657273838145435,240:1.3163963725085925):0.151402686975042,(148:0.8502621462463897,((187:0.2218914614395655,227:0.2218914614395655):0.1875283161639416,((203:0.0143421338853984,189:0.0143421338853984):0.25115917757000794,234:0.26550131145540634):0.14391846614810078):0.4408423686428826):0.6175369132372448):0.07507504401318077,(116:0.21804866081771657,186:0.21804866081771657):1.3248254426790989):0.10524927559301678,(108:0.5575308819625084,90:0.5575308819625084):1.0905924971273238):0.08457795190226658,(((((46:0.25823318903568193,99:0.25823318903568193):0.6298639004154232,115:0.8880970894511051):0.00352927147474813,(((((237:0.20741715038729058,36:0.20741715038729058):0.18443102512754628,212:0.39184817551483686):0.0032976040457390887,(124:0.23283126437516266,161:0.23283126437516266):0.1623145151854133):0.011369160272506118,9:0.40651493983308207):0.1494414444023171,220:0.5559563842353992):0.33566997669045406):0.3165344934937879,118:1.2081608544196412):0.3617909098173575,((((((162:0.034945690780505245,132:0.034945690780505245):0.19242962254313856,42:0.2273753133236438):0.17549828083562113,(((164:0.02493672768545574,60:0.02493672768545574):0.14923286629530885,(31:2.861517250267731E-4,207:2.861517250267731E-4):0.17388344225573782):0.19993388204264573,(71:0.22184743146580654,(157:0.050408750672636504,80:0.050408750672636504):0.17143868079317004):0.15225604455760378):0.028770118135854617):0.23343520726585498,154:0.6363088014251199):0.2880502604352615,(((111:0.03219949787489451,94:0.03219949787489451):0.17503222233260063,(199:0.017591236215728223,104:0.017591236215728223):0.18964048399176692):0.040289566553747935,141:0.24752128676124308):0.6768377750991383):0.355313383051572,84:1.2796724449119534):0.29027931932504525):0.1627495667551):0.1616741410359196,(214:1.5211513817045366,((86:0.6423759217102636,((236:0.3322929367219407,185:0.3322929367219407):0.1827825144107429,231:0.5150754511326836):0.12730047057757998):0.73511664196663,((180:0.3862334779611709,158:0.3862334779611709):0.5225168133863127,(54:0.0995536332436231,119:0.0995536332436231):0.8091966581038605):0.4687422723294099):0.1436588180276429):0.3732240903234818):0.22896790267142003,((((((((((48:0.4241004269162314,243:0.4241004269162314):0.09784151268971453,((74:0.04626858886278473,45:0.04626858886278473):0.3274362112137248,92:0.3737048000765095):0.1482371395294364):0.5175644010452773,225:1.0395063406512233):0.026114645996676566,(((83:0.1946664410900527,112:0.1946664410900527):0.2158928831153264,238:0.4105593242053791):0.41494843548179916,((95:0.6624421630615618,(98:0.3736316787398084,63:0.3736316787398084):0.2888104843217534):0.015930502252089696,47:0.6783726653136515):0.1471350943735268):0.24011322696072157):0.031243061705280573,73:1.0968640483531804):0.3638019883363919,((((160:0.11211745172100596,213:0.11211745172100596):0.16150366590205723,197:0.2736211176230632):0.4628560840370626,33:0.7364772016601258):0.6701138844558105,(((120:0.4487509591897667,204:0.4487509591897667):0.23514115314826545,((((117:0.23445717567867108,((137:0.056331038895975016,40:0.056331038895975016):0.12404803586902835,(139:0.10310711691621854,122:0.10310711691621854):0.07727195784878482):0.054078100913667715):0.027281350236161295,191:0.2617385259148324):0.2266784992072739,(((107:0.05571036530712714,93:0.05571036530712714):0.12217386013969955,205:0.1778842254468267):0.061749174821942976,72:0.23963340026876967):0.2487836248533366):0.024949902744727748,(64:0.3232306519751287,66:0.3232306519751287):0.19013627589170534):0.17052518447119813):0.6826161111792608,(((134:1.0617729361589328,(3:0.369845068592098,133:0.369845068592098):0.6919278675668348):0.06500744366185152,(((((232:0.021546584301352034,244:0.021546584301352034):0.344142745051379,59:0.365689329352731):0.08940911502227822,(51:0.06065352159978854,155:0.06065352159978854):0.3944449227752207):0.44485531898928143,(200:0.17317170680705862,169:0.17317170680705862):0.726782056557232):0.12099209756631146,((145:0.6541678249491143,(125:0.6113229366929425,18:0.6113229366929425):0.04284488825617183):0.307507747966552,(((173:0.18571664984559622,109:0.18571664984559622):0.17867471364482457,(7:0.0658969433606229,215:0.0658969433606229):0.2984944201297979):0.5009386486589174,85:0.8653300121493381):0.09634556076632816):0.059270288014935835):0.10583451889018214):0.12169744517167924,((((70:0.24626376861760324,68:0.24626376861760324):0.202113365316825,(194:0.2759212994177287,129:0.2759212994177287):0.17245583451669955):0.1082402998360783,57:0.5566174337705065):0.5724768165144474,((44:0.8676844249624589,(81:0.048359004317534815,223:0.048359004317534815):0.8193254206449241):0.03144437451077775,(((((135:0.02217996438076053,37:0.02217996438076053):0.20079958631531447,15:0.222979550696075):0.30733208388215183,210:0.5303116345782268):0.019679586316499353,241:0.5499912208947262):0.10025402948977513,52:0.6502452503845013):0.24888354908873533):0.22996545081171726):0.1193835747075096):0.11803039852482944):0.040082862598643354):0.05407495057363598):0.0074844678674776954,((((230:0.20357841368630236,50:0.20357841368630236):0.34349556416463,(178:0.012266430667138728,96:0.012266430667138728):0.5348075471837936):0.7625820889503805,((28:0.9434330370416533,(((61:0.07242909393280295,23:0.07242909393280295):0.23479701936494912,67:0.30722611329775207):0.11523249249406309,100:0.42245860579181516):0.5209744312498381):0.32781197583206234,5:1.2712450128737156):0.03841105392759725):0.0033997287112403907,(((235:0.02789630554792666,103:0.02789630554792666):0.10076698096962167,6:0.12866328651754833):0.8103457782957728,((21:0.05467586259720836,127:0.05467586259720836):0.7499656993517843,(174:0.7656118147305466,(142:0.5850056802792643,102:0.5850056802792643):0.18060613445128237):0.03902974721844599):0.13436750286432853):0.3740467306992321):0.15509470904449674):0.038425441676050665,((136:0.9677314899312848,(((146:0.011349922960334435,239:0.011349922960334435):0.2842512119673697,(201:0.11182495242383439,177:0.11182495242383439):0.18377618250386973):0.3129454953546442,(143:0.0035453008061563374,216:0.0035453008061563374):0.605001329476192):0.3591848596489364):0.2348247769399423,(((((149:0.22758379023380604,166:0.22758379023380604):0.441688397673925,((41:0.3322776423307512,17:0.3322776423307512):0.14240197672737986,((152:0.18721514171821152,159:0.18721514171821152):0.06417088823374106,((226:0.15128689471622314,179:0.15128689471622314):0.0067865726840588,167:0.15807346740028194):0.09331256255167064):0.22329358910617847):0.19459256884959997):0.0010930897008205598,(((((130:0.13697614690031212,19:0.13697614690031212):0.1556390440223976,(221:0.10231532304161384,113:0.10231532304161384):0.19029986788109587):0.18149349857988994,(110:0.18834896760392006,20:0.18834896760392006):0.2857597218986796):0.003475660991485885,(14:0.07492552776593131,217:0.07492552776593131):0.40265882272815423):0.05115386060759408,(53:0.35364413984485843,140:0.35364413984485843):0.1750940712568212):0.14162706650687196):0.38161461281194176,245:1.0519798904204933):0.09874570459598808,((156:0.9507599951088002,((121:0.004259389439384087,123:0.004259389439384087):0.06344738090632385,39:0.06770677034570793):0.8830532247630922):0.19462081399941344,(((((114:0.47470230915442047,34:0.47470230915442047):0.16424582016417677,((65:0.22392905261685803,105:0.22392905261685803):0.3909817191951903,(((29:0.0791638720399086,69:0.0791638720399086):0.039560932068283705,181:0.1187248041081923):0.28923515555720103,208:0.40795995966539333):0.206950812146655):0.024037357506548895):0.23532349664389995,((106:0.0015754429076402232,43:0.0015754429076402232):0.24118348050045535,126:0.24275892340809557):0.6315127025544016):0.08721296138400114,(((2:0.0304985617330793,4:0.0304985617330793):0.00863604774277027,168:0.03913460947584957):0.8163217477494533,(175:0.20482219199414375,76:0.20482219199414375):0.6506341652311591):0.10602823012119544):0.10346825662582781,((55:0.3951313616811536,8:0.3951313616811536):0.665434039630213,((((58:0.6203790340924247,(88:0.3838148173248017,218:0.3838148173248017):0.23656421676762296):0.0307294654231558,171:0.6511084995155805):0.04984065823744177,195:0.7009491577530222):0.24342106777978745,((151:0.21100802668854302,49:0.21100802668854302):0.600667885920686,16:0.811675912609229):0.1326943129235807):0.11619517577855687):0.004387442660959584):0.08042796513588746):0.005344785908267813):0.05183067185474566):0.3040196793618736):0.17026942978592519,((((27:0.008745783438392607,176:0.008745783438392607):0.0346849214713818,82:0.04343070490977441):0.2166861390994037,193:0.2601168440091781):1.166625135381603,((12:0.2814908996940111,97:0.2814908996940111):0.8716791748539319,(25:0.41496509775598067,((188:0.018499592191076886,138:0.018499592191076886):0.25668984144183504,89:0.27518943363291193):0.13977566412306874):0.7382049767919623):0.2735719048428382):0.25010339662824466):0.3913503628622723,(((((182:0.2836112687761534,242:0.2836112687761534):0.08505667133173889,11:0.3686679401078923):0.21699666682762087,131:0.5856646069355131):0.1867452643800538,(((153:0.10983208776720721,228:0.10983208776720721):0.20167938981497846,211:0.31151147758218567):0.146907349821229,((24:0.07463738743032255,190:0.07463738743032255):0.3750026945181548,147:0.4496400819484774):0.008778745454937287):0.31399104391215227):0.924967169099673,((183:0.1626887650299671,87:0.1626887650299671):1.3028585412048237,((((192:1.0086093884575675,(184:0.9581405521641553,163:0.9581405521641553):0.050468836293412256):0.2812539220150916,(((((75:0.5980640452438888,165:0.5980640452438888):0.4282395110734112,78:1.0263035563173):0.10997888368027486,((((144:0.07076058964647514,13:0.07076058964647514):0.21414114144085605,128:0.2849017310873312):0.38482005694835086,26:0.669721788035682):0.35842786677739524,170:1.0281496548130773):0.10813278518449754):0.01609087084087868,(22:0.025400280250206908,1:0.025400280250206908):1.1269730305882466):0.12421314506710956,(((91:0.49465986133119433,206:0.49465986133119433):0.28429416551135644,62:0.7789540268425508):0.34337813526864536,(209:0.7202780440806722,(219:0.06423486056776584,56:0.06423486056776584):0.6560431835129064):0.40205411803052393):0.15425429379436695):0.013276854567096041):0.1291352770611054,((10:0.114182440439905,222:0.114182440439905):0.8648036753645074,150:0.9789861158044124):0.4400124717293521):0.03990426993392382,(202:0.920512328860912,77:0.920512328860912):0.5383905286067763):0.006644448767102462):0.23182973418044917):0.3708186984660582):0.05514763581814014);";
//        tips[2] = 245;

        Tree tree;

        double logLvolzSIR2009;

        for (int j = 0; j < trees.length; j++) {
            for (int p = 0; p < 3; p++) {
                
                writer = new PrintStream("vary_" + param_name[p] + "_" + j + ".txt");
                writer.println("#tree = " + trees[j]);
                writer.println("#tips = " + tips[j]);

                tree = new TreeParser(trees[j], false);

                TreeIntervals treeIntervals = new TreeIntervals();
                treeIntervals.init(tree);
                
                System.arraycopy(paramTruth, 0, paramCurrent, 0, 3);
                double paramStep = (paramUpperBound[p] - paramLowerBound[p]) / nsteps;

                writer.format("%s\tvolzSIR\tvolz2009\n", param_name[p]);

                for (int i = 1; i <= nsteps; i++) {

                    paramCurrent[p] = paramLowerBound[p] + i * paramStep;
                    writer.print(paramCurrent[p] + "\t");

                    
                    // Volz 2012
                    
                    DeterministicCoalescentSIR volzSIR = new DeterministicCoalescentSIR();
                    volzSIR.setInputValue("beta", new RealParameter(String.valueOf(paramCurrent[0])));
                    volzSIR.setInputValue("gamma", new RealParameter(String.valueOf(paramCurrent[1])));
                    volzSIR.setInputValue("n_S0", new RealParameter(String.valueOf(paramCurrent[2])));
                    volzSIR.setInputValue("origin", new RealParameter("13.020168"));
                    volzSIR.setInputValue("integrationStepCount", 1000);
                    volzSIR.initAndValidate();
                    
                    Volz2009TreeDistribution coal = new Volz2009TreeDistribution();

                    // "Restructured since Denise wrote this" -Alex
                    //Volz2012PopulationFunction volz2012 = new Volz2012PopulationFunction();

                    coal.initByName("treeIntervals", treeIntervals,
                            "volzSIR", volzSIR);
                    
                    double logLvolzSIR = coal.calculateLogP();
                    writer.print((Double.isInfinite(logLvolzSIR) ? "NA" : logLvolzSIR) + "\t");
                    
                    
                    // Volz 2009
                    // "Restructured since Denise wrote this" -Alex
                    
                    Volz2009TreeDistribution volzSIR2009 = new Volz2009TreeDistribution();
                    volzSIR2009.setInputValue("treeIntervals", new TreeIntervals(tree));
                    volzSIR2009.setInputValue("beta", new RealParameter(String.valueOf(paramCurrent[0])));
                    volzSIR2009.setInputValue("delta", new RealParameter(String.valueOf(paramCurrent[1])));
                    volzSIR2009.setInputValue("S0", new RealParameter(String.valueOf(paramCurrent[2])));

                    try {
                        volzSIR2009.initAndValidate();
                        logLvolzSIR2009 = volzSIR2009.calculateLogP();

                    } catch (Exception e) {
                        logLvolzSIR2009 = Double.NaN;
                    }

                    writer.print((Double.isInfinite(logLvolzSIR2009) ? "NA" : logLvolzSIR2009) + "\t");


                    // BDSIR

//                    HybridSEIREpidemic sir = new HybridSEIREpidemic();
//                    sir.initByName(
//                            "tree", tree,
//                            "birth", new RealParameter(String.valueOf(paramCurrent[0])),
//                            "death", new RealParameter(String.valueOf(paramCurrent[1])),
//                            "sampling", "0",
//                            "origin", new RealParameter("2.5"),
//                            "S0", new RealParameter(String.valueOf(paramCurrent[2])),
//                            "Nt", "10001",
//                            "Nsamples", "101",
//                            "simulationType", "SAL", "isDeterministic", false);

//                    sir.initTraj(paramCurrent[0], 0., new Double[]{paramCurrent[1]}, new Double[]{0.}, 2.5, tree.getLeafNodeCount(), 100, 1000, sir.times);

//                    BDSIR bdsir = new BDSIR();
//                    bdsir.setInputValue("tree", tree);
//                    bdsir.setInputValue("birthRate", new RealParameter(new Double[]{paramCurrent[0] * paramCurrent[2]}));
//                    bdsir.setInputValue("deathRate", new RealParameter(String.valueOf(paramCurrent[1])));
//                    bdsir.setInputValue("samplingRate", new RealParameter("0"));
//                    bdsir.setInputValue("rho", new RealParameter(new Double[]{tips[j] / 245.}));
//                    bdsir.setInputValue("S0", new RealParameter(String.valueOf(paramCurrent[2])));
//                    bdsir.setInputValue("origin", new RealParameter("2.5"));
//                    bdsir.setInputValue("dS", new RealParameter(sir.get_dS()));
//                    bdsir.setInputValue("dR", new RealParameter(sir.get_dR()));
//                    bdsir.setInputValue("checkTreeConsistent", false);

//                    bdsir.initAndValidate();


//                    double logLbdsir;
//                    try {
//                        logLbdsir = bdsir.calculateTreeLogLikelihood(tree);
//
//                    } catch (Exception e) {
//                        logLbdsir = Double.NaN;
//                    }

//                    writer.println((Double.isInfinite(logLbdsir) ? "NA" : logLbdsir) + "\t");

                }
                writer.flush();
                writer.close();
            }

        }
    }

}
